# stuff that gets performed every time this module is imported (just once though)
from numpy import *
from genutils.io import *
from genutils.special_functions import *
import random
import cPickle
import os
import copy

# seed the random number generator FIXME this should be in the config file
#random.seed(41)

class Cpopulation:
    """ This class holds many cells and exercises them and 'evolves' """
    def __init__(self,config_file="configure"):
        self.cells = []
        self.scores = []
        self.avg_scores = []
        self.generation_number = 0
        self.num_runs = int(from_config("generations",config_file=config_file))

    def exercise(self):
        # exercise all the cells and archive them
        self.generation_number += 1

        print "Running GA on %d cells in generation %d" % (len(self.cells),self.generation_number)
        cell_number = 0
        self.scores = []

        for cell in self.cells:
            cell_number += 1
            # no need to test the number one cell (eliteism)
            if cell_number != 1:
                cell.exercise()
            self.scores.append(cell.score)
            print "     ---> cell %d had a score of %f" % (cell_number,cell.score)
            dir_name = "../../data/%d-Hz/generation-%04d/%05d-%06d" % (cell.NEURON_frequency,self.generation_number,\
                                                                            cell.score,random.randint(0,999999))
            back_to_code_path = "../../../../code/"
            cell.archive(dir_name,back_to_code_path)
            print "     ===> cell has been archived to " + dir_name
        # record the average scores of the generations
        self.avg_score =  sum(self.scores)/len(self.scores)
        self.avg_scores.append(self.avg_score)
        print "   ~~~ generation completed with average score of %f" % (self.avg_scores[-1])

    def evolve(self,cross_rate=0.08,mutate_rate=0.01):

        sorted_scores = copy.deepcopy(self.scores)
        sorted_scores.sort()

        high_score = sorted_scores[-1]
        high_cell = self.cells[self.scores.index(high_score)]
        new_cells = []

        new_cells.append(high_cell) # elitism

        # make a set of cells to procreate
        probs = copy.deepcopy(self.scores)
        tot_score = sum(self.scores)
        for i in range(len(probs)):
            probs[i] /= tot_score
        pro_cells = []
        # start adding unaltered cels to the pro_cells list
        for i in range(len(self.cells)-1):
            r = random.uniform(0,1)
            cell_num = 0
            cell_prob = probs[0]
            while r > cell_prob and cell_num+1 <= len(self.cells):
                cell_num += 1
                cell_prob += probs[cell_num]
            # we've selected the cell index with cell_num so let's add it to pro_cells
            pro_cells.append(copy.deepcopy( self.cells[cell_num] ))

        # cross and then mutate the pro_cells with one another
        for i in range(len(pro_cells)/2):
            pro_cells[2*i].chromosome.cross_with(pro_cells[2*i+1].chromosome,cross_rate)
        for pro_cell in pro_cells:
            pro_cell.chromosome.mutate(mutate_rate)

        # finally, add the procreated cells to the new_cell list
        new_cells.extend(pro_cells)
        self.cells = new_cells
        


class Ccell:
    """ This top-level class holds two trees and has some methods
            for running simulations using NEURON scripts
-----------------------------------------------------------------"""
    def __init__(self,config_file="configure"):
        # just randomly sample from the configure file making a cell 
        # with trees sharing the same genes.
        self.config_file = config_file
        self.ensemble_size = int(from_config("ensemble_size",config_file=config_file))
        self.max_synapses = int(from_config("max_synapses_per_tree",config_file=config_file))
        self.NEURON_tstop = from_config("simulation_length", config_file=config_file)
        self.NEURON_dt = from_config("simulation_time_step_size", config_file=config_file)
        self.NEURON_sfr = from_config("simulation_synapse_firing_rate", config_file=config_file)
        self.NEURON_frequency = from_config("simulation_input_frequency",config_file=config_file)
        self.NEURON_stc = from_config("simulation_synapse_time_constant",config_file=config_file)
        self.NEURON_smc = from_config("simulation_synapse_max_conductance",config_file=config_file)
        self.NEURON_srp = from_config("simulation_synapse_reversal_potential",config_file=config_file)
        self.NEURON_numruns = from_config("simulation_numruns",config_file=config_file)
        self.NEURON_fullphase = from_config("simulation_fullphase",config_file=config_file)
        self.chromosome = Cchromosome(config_file=config_file)
        self.trees = []
        for i in range(self.chromosome.num_genes):
            self.trees.append( Ctree(self.chromosome.genes[i]) ) # construct the tree object instances with the specified gene
        self.generate()

        # impose sanity checks
        while self.trees[0].num_synapses > self.max_synapses:
            self.chromosome.lower_ls_density(config_file=config_file)
            self.trees = []
            for i in range(self.chromosome.num_genes):
                self.trees.append( Ctree(self.chromosome.genes[i]) ) # construct the tree object instances with the specified gene
            self.generate()

    def generate(self):
        # create the soma
        soma = self.soma = Ctree_segment()
        soma.ID = -1
        soma.length = from_config("soma_length",config_file=self.config_file)
        soma.nsegs = 3
        soma.end_diameter = soma.diameter =  from_config("soma_diameter",config_file=self.config_file)
        soma.beg_3d = array([0,soma.length/2,0])
        soma.end_3d = array([0,-soma.length/2,0])

        # create the trees
        self.ID_counter = 0
        # tree 1
        self.ID_counter = self.trees[0].generate(ID_counter=self.ID_counter,begining=soma.beg_3d,soma=soma)
        self.trees[0].add_synapse_locations()
        # tree 2
        self.ID_counter = 0
        self.ID_counter = self.trees[1].generate(ID_counter=self.ID_counter,begining=soma.end_3d,soma=soma)
        self.trees[1].add_synapse_locations()

    def gnuplot_print(self,ofile):
        self.soma.gnuplot_print(ofile,draw_axes=True)
        self.trees[0].gnuplot_print(ofile,draw_synapses="left")
        self.trees[1].gnuplot_print(ofile,draw_synapses="right")

    def neuron_print(self,ofile):
        # print simulation settings
        ostring = '%f\t%f\t%f\t%f\t%f\t%f\t%f\t%d\t%f\n' % (self.NEURON_tstop,self.NEURON_dt,self.NEURON_sfr,\
                                    self.NEURON_frequency,self.NEURON_stc,\
                                    self.NEURON_smc,self.NEURON_srp,self.NEURON_numruns,self.NEURON_fullphase)
        ofile.write(ostring)

        # print the number of lines to read
        ofile.write(str(len(self.trees))+"\n")
        for tree in self.trees:
            tree.neuron_print_num_lines(ofile)
        for tree in self.trees:
            tree.synapse_print_num_lines(ofile)

        # print soma information
        s = self.soma # to make things shorter to type
        ostring = '%d\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\n' % (s.ID,\
                s.beg_3d[0],s.beg_3d[1],s.beg_3d[2],s.diameter,\
                s.end_3d[0],s.end_3d[1],s.end_3d[2],s.end_diameter,s.nsegs)
        ofile.write(ostring)

        # print the trees information
        for tree in self.trees:
            # print the segments geometric information
            tree.neuron_print(ofile)
        for tree in self.trees:
            # print the synapse locations
            tree.synapse_print(ofile)
        # death to fascism!

    def archive(self,dir_name,get_back_to_code_path):
        # dir_name = self.dir_name = "../../data/%d-Hz/%05d-%06d" % (self.NEURON_frequency,self.score,random.randint(0,999999))
        # get_back_to_code_path = '../../../code/'

        self.gnuplot_print(open('plot3d.out','w'))
        # pickle the cell
        cPickle.dump(self,open('cell.cPickle','w'))
        # make a directory to put all the stuff into
        
        os.system("mkdir -p " + dir_name) # -p to not yell about intermediate directories
        # copy everything into the new directory
        os.system("cp -f *.out* " + dir_name)
        os.system("cp -f configure " + dir_name)
        os.system("cp -f cell.cPickle " + dir_name)

        # make symbolic links to python and gnuplot scripts
        for file in os.listdir('..'): # one directory up from the run directory
            if file.find('.gnuplot') != -1:
                os.system('ln -s ' + get_back_to_code_path + file + ' ' + dir_name + '/' + file)
            if file.find('.py') != -1:
                os.system('ln -s ' + get_back_to_code_path + file + ' ' + dir_name + '/' + file)
            if file.find('.hoc') != -1:
                os.system('ln -s ' + get_back_to_code_path + file + ' ' + dir_name + '/' + file)
        
        



    def exercise(self):
        # generate an ensemble of realizations of the cell and use NEURON to give it a fitness value
        self.sim_results = []
        for i in range(self.ensemble_size):
            self.generate() # generate a new instance of this cell with the same chromosome
            self.neuron_print(open('neuron.out','w')) # print the instructions for NEURON to follow
            os.system('nrniv exercise_cell.hoc >& sim%d.out' % (i+1) ) # run the simulation
            # save the results of each simulation run
            self.sim_results.append(read_dict_float('results.out'))
        # calculate the cell's score
        firing_rates = []
        avg_results = self.sim_results[0].copy()
        for key in self.sim_results[0].keys():
            for i in range(len(self.sim_results))[1:]:
                 avg_results[key] += self.sim_results[i][key]
        for key in avg_results.keys():
            avg_results[key] /= self.ensemble_size    
            firing_rates.append(avg_results[key])
        firing_rates.sort()
        self.score = firing_rates[-1] - firing_rates[0]
            
        

class Ctree:
    """ This class is designed to represent a whole dendritic tree.

    Members:
        It has a list of all [segments] in the tree
        It has a pointer to the [gene] that describes its generation.
        It has a pointer to the [root] segment.
        It has a [longest_path] distance from soma.
        It has a [total_length] (sum all segments end to end)
    Methods:
        generate(ID_counter=0, flat_tree=False, steps=-1)
        add_synapse_locations()
        synapse_print(ofile)
        diameter_vs_dist_plot(ofile):
        gnuplot_print(ofile, draw_synapses=False):
        neuron_print(ofile):
-------------------------------------------------------------------"""
    def __init__(self,gene):
        # static initialization
        self.longest_path = 0
        self.total_length = 0
        self.esampled_length = [] # USED FOR PLOTTING ONLY
        self.esampled_asymmetry = [] # USED FOR PLOTTING ONLY
        self.esampled_branching_angle = [] # USED FOR PLOTTING ONLY
        self.esampled_ls_density = [] # USED FOR PLOTTING ONLY

        # variable initialization
        self.gene = gene
        self.no_branching_length = gene.no_branching_length

    def generate(self, ID_counter=0,flat_tree=False, steps=-1, begining=array([0,0,0]),soma=False):
        # this proceedure generates all the segments in the tree based on
        #    the gene describing this tree.
        # steps specifies the number of times we let the generation loop run
        #    it's used for making some plots but is generally of limited utility.
        # steps=-1 (default) will run a full generation process

        # initialize some things
        gene = self.gene
        self.longest_path = 0
        self.total_length = 0
        g_stack = [] # growing segments stack
        self.segments = [] # list of all created segments
        self.sampled_branching_angle = [] # USED FOR PLOTTING ONLY
        self.sampled_asymmetry = [] # USED FOR PLOTTING ONLY
        self.sampled_length = [] # USED FOR PLOTTING ONLY
        self.sampled_ls_density = [] # USED FOR PLOTTING ONLY

        # ---------------------- ROOT -------------------------------------------------
        self.sampled_start_diameter = gene.sample_start_diameter()
        root_length = gene.sample_length(0)
        self.sampled_length.append([0,root_length]) # USED FOR PLOTTING ONLY
        root = self.root = Ctree_segment(is_root=True,diameter=self.sampled_start_diameter,\
                                         theta=gene.root_theta,phi=gene.root_phi,\
                                         length=root_length )
        # make the root it's own parent (just so everyone has a parent) unless the soma was passed
        if soma:
            root.parent = soma
        else:
            root.parent = root

        # set the root's ID and update the ID counter
        root.set_ID(ID_counter)
        ID_counter += 1

        # set the diameter tapering rate for the tree
        gene.taper_rate = root.diameter/self.no_branching_length

        # set the end diameter and adjust the length if necessary
        root.set_end_diameter(gene.taper_rate)
        root.set_nseg(gene.nseg_per_length)
        root.beg_3d = begining
        root.calc_end_3d()

        # grow the rest of the tree
        g_stack.append(self.root)
        self.segments.append(self.root)
        #-------------------- MAIN GROWING LOOP ------------------------------------------------
        while len(g_stack) and steps != 0:
            steps -= 1
            # make the first element of the generation stack the new parent and generate its children
            parent = g_stack.pop()
            # D is the distance to the soma of the children we're going to generate
            D = parent.dist_to_soma + parent.length
            #--------------------- SAMPLE PARAMETERS -----------------------------------------
            # we need to know the children's branching angle and asymmetry parameter to make them
            branching_angle = gene.sample_branching_angle(D)
            asymmetry = gene.sample_asymmetry(D)
            self.sampled_branching_angle.append([D,branching_angle]) # USED FOR PLOTTING ONLY
            self.sampled_asymmetry.append([D,asymmetry]) # USED FOR PLOTTING ONLY

            # asymmetry 1 and 2 used to determine angles and diameters
            asymmetry_1 = 1/(asymmetry + 2)
            asymmetry_2 = 1 - asymmetry_1
            # determine the angles of the children and their starting diameters
            #   using Rall's power law with exponent 1.5 to provide impedence matching
            #   d_p**1.5 = d_c1**1.5 + d_c2**1.5
            # ------------------ EQUATIONS FOR DETERMINING DIAMETERS -----------------------
            # d_c1 * asymmetry_1 = d_c2 * asymmetry_2
            # d_c1**1.5 * (1 + (asymmetry_1 / asymmetry_2)**1.5) = d_p**1.5
            # d_c1 =( d_p**1.5 / (1 + (asymmetry_1 / asymmetry_2)**1.5) )**(1/1.5)

            # -------------------- CHILD 1 -------------------------------------------------
            # create the child object instance
            child1 = Ctree_segment.with_parent(parent)

            # set the child's ID and update the ID counter
            child1.set_ID(ID_counter)
            ID_counter += 1

            # set the diameter and angles
            child1.diameter = (  parent.end_diameter**1.5 / \
                              (1 + (asymmetry_1 / asymmetry_2)**1.5) )**(1/1.5)
            child1.phi = branching_angle*asymmetry_1
            # set the theta angle for the first child randomly
            #     unless we want a flat tree
            if flat_tree:
                child1.theta = 0.0
            else:
                child1.theta = random.random() * 2 * pi

            # set the length
            child1.length = gene.sample_length(D)
            self.sampled_length.append([D,child1.length]) # USED FOR PLOTTING ONLY

            # set the end diameter (and adjust length down if necessary)
            child1.set_end_diameter(gene.taper_rate)
            child1.set_nseg(gene.nseg_per_length)

            # determine the 3d cartesian coordinates of the segment's endpoint
            child1.calc_end_3d()

            # -------------------- CHILD 2 -------------------------------------------------
            # create the child object instance
            child2 = Ctree_segment.with_parent(parent)

            # set the child's ID and update the ID counter
            child2.set_ID(ID_counter)
            ID_counter += 1

            # set the diameter and angles
            child2.diameter = (  parent.end_diameter**1.5 / \
                              (1 + (asymmetry_2 / asymmetry_1)**1.5) )**(1/1.5)
            child2.phi = branching_angle*asymmetry_2
            child2.theta = child1.theta + pi # second child theta just 180 degrees from first

            # set the length
            child2.length = gene.sample_length(D)
            self.sampled_length.append([D,child2.length]) # USED FOR PLOTTING ONLY

            # set the end diameter (and adjust length down if necessary)
            child2.set_end_diameter(gene.taper_rate)
            child2.set_nseg(gene.nseg_per_length)

            # determine the 3d cartesian coordinates of the segment's endpoint
            child2.calc_end_3d()

            # ------------------- KEEP TRACK OF CHILDREN ------------------------------------------
            # segments is a list of all tree segments in this tree
            self.segments.append(child1)
            self.segments.append(child2)

            # if the child isn't a leaf, add it to the generation stack
            if child1.end_diameter > 0:
                g_stack.append(child1)
            if child2.end_diameter > 0:
                g_stack.append(child2)
        #-------------------- END MAIN GROWING LOOP ------------------------------------------------

        # now that the tree is built, let's determine its overall length and stuff
        for segment in self.segments:
            # total length
            self.total_length += segment.length
            # longest path
            self.longest_path = max(self.longest_path, segment.dist_to_soma + segment.length) 
        # keep track of all sampled variables ever sampled by this tree
        self.esampled_length.extend(self.sampled_length) # USED FOR PLOTTING ONLY
        self.esampled_asymmetry.extend(self.sampled_asymmetry) # USED FOR PLOTTING ONLY
        self.esampled_branching_angle.extend(self.sampled_branching_angle) # USED FOR PLOTTING ONLY
        # return the ID_counter so we can have unique ID's for every segment in multiple trees
        return ID_counter

    def add_synapse_locations(self):
        # add synapse locations to the tree based on the gene's ls_density
        self.num_synapses = 0
        self.sampled_ls_density = []
        for segment in self.segments:
            keep_adding = True
            dl = self.longest_path/100.0 # we'll assume that the ls_density is constant over intervals of this distance
            D = segment.dist_to_soma
            # make sure segment is at least dl long
            if segment.length < dl:
                # set dl to be the length of the segment
                dl = segment.length
                keep_adding = False
            # determine the number of synapses to put in this part of the segment
            ls_density = self.gene.sample_ls_density(D)
            self.sampled_ls_density.append([D,ls_density])
            if random.uniform(0,1) <= ls_density * dl:
                segment.add_synapse_location(D)  
                self.num_synapses += 1
            while (D+dl < segment.dist_to_soma + segment.length) and keep_adding:
                D += dl
                # make sure there is at least dl length left in the segment
                if  D + dl > segment.dist_to_soma + segment.length:
                    # set dl to be the length to the end of the segment
                    dl = segment.dist_to_soma + segment.length - D
                    keep_adding = False
                # determine the number of synapses to put in this part of the segment
                ls_density = self.gene.sample_ls_density(D)
                self.sampled_ls_density.append([D,ls_density])
                if random.uniform(0,1) <= ls_density * dl:
                    segment.add_synapse_location(D)  
                    self.num_synapses += 1
        self.esampled_ls_density.extend(self.sampled_ls_density)


    def calc_avg_synapse_distance(self):
        # calculate the average synapse distance from the soma
        dist = 0
        for segment in self.segments:
            for syn in segment.synapses:
                dist += syn*segment.length + segment.dist_to_soma
        self.avg_synapse_distance = dist/self.num_synapses

    def synapse_print_num_lines(self,ofile):
        # print out how many synapses there are in this tree
        ofile.write(str(self.num_synapses) + "\n")
        

    def synapse_print(self,ofile):
        # print out all the synapse locations and associated segment IDs
        # print out the synapse locations and the associated segment ID number
        for segment in self.segments:
            for synapse_location in segment.synapses:
                ostring = '%d\t%f\n' % (segment.ID,synapse_location)
                ofile.write(ostring)
                
        
    def diameter_vs_dist_plot(self, ofile):
        t = "\t"
        n = "\n"
        for segment in self.segments:
            D = str(segment.dist_to_soma)
            diam = str(segment.diameter)
            outstring = D + t + diam + n
            ofile.write(outstring)
            
    def gnuplot_print(self, ofile, draw_synapses=False):
        self.root.gnuplot_print(ofile,draw_axes=1,draw_synapses=draw_synapses)
        for segment in self.segments:
            segment.gnuplot_print(ofile,draw_synapses=draw_synapses)

    def neuron_print_num_lines(self,ofile):
        ofile.write( str(len(self.segments)) + "\n" )
        
    def neuron_print(self, ofile):
        for segment in self.segments:
            segment.neuron_print(ofile)

class Cgene:
    """ This class holds all the information to build one dendritic tree

    Members:
        It has a list of allyls which each help describe a dendritic tree
    Methods:
        calc_dist_norm()
        gnuplot_print(ofile)
        parameters_print(ofile)
        list_form(parameter_name, num_pts=200, error='no/upper/lower'):

        sample_start_diameter(self):
        sample_branching_angle(dist_to_soma):
        sample_length(dist_to_soma):
        sample_asymmetry(dist_to_soma):
        sample_ls_density(dist_to_soma): # the linear synaptic density
-------------------------------------------------------------------"""

    def __init__(self, dict, dist_norm=False):
        self.__dict__.update(dict) # set variables based on the passed dictionary object
        self.nseg_per_length = 200.00/self.no_branching_length
        # variables which are functions of a distance to soma are described by parameters
        #     which represent gaussians on a normalized length scale.  This scale has to
        #     be calculated.  Once it has been calculated 0 represents the root's base and
        #     1 is the maximal distance to soma for any segment in a tree.
        if dist_norm:
            self.dist_norm = dist_norm
        else:
            self.dist_norm = self.no_branching_length 
            self.calc_dist_norm()

    def calc_dist_norm(self):
        # make a tree (or ten) and determine the dist_norm iteratively
    
        self.nseg_per_length = 200.00/self.no_branching_length
        tree = Ctree(self)
        tree.generate()
        # while the error is greater than 10% keep working
        while (tree.longest_path - self.dist_norm)**2 > 0.1**2:
            tree.generate()
            # if the normalizing factor is too big, cut it down
            if tree.longest_path < self.dist_norm:
                # subtract half the difference
                self.dist_norm -= (self.dist_norm - tree.longest_path) * 0.85
            else:
                break


    def gnuplot_print(self, ofile):
        for i in range(1000):
            t = "\t"
            n = "\n"
            D = (float(i)/1000)*self.dist_norm
            a = str( self.sample_asymmetry(D) ) + t
            b = str( self.sample_branching_angle(D) ) + t
            c = str( self.sample_length(D) ) + n
            output_string = str(D) + t + a + b + c 
            ofile.write(output_string) 

    def parameters_print(self, ofile):
        s = " =  "
        n = "\n"
        ofile.write("start_diameter" + s + str(self.start_diameter) + n)
        ofile.write("no_branching_length" + s + str(self.no_branching_length) + n)
        ofile.write("parameter_noise_sd" + s + str(self.parameter_noise_sd) + "%" + n)
        

    def sample_start_diameter(self):
        # return the root's starting diameter
        return random.gauss(self.start_diameter,\
                            self.start_diameter*self.parameter_noise_sd/100)

#    def sample_no_branching_length(self):
#        # return the root's starting diameter
#        return random.gauss(self.no_branching_length,\
#                            self.no_branching_length*self.parameter_noise_sd/100)

    def sample_ls_density(self,dist_to_soma): # the linear synaptic density
        D = dist_to_soma/self.dist_norm
        ls_density = gaussian(self.ls_density_center,self.ls_density_peak,self.ls_density_width,D)
        ls_density += self.ls_density_baseline
        # no noise associated with the linear synaptic density since synapse position is determined
        #   stochastically anyways
        return ls_density

    def sample_asymmetry(self,dist_to_soma):
        # return the branching asymmetry value given its distance from the soma.
        #
        # translate <dist_to_soma> to normalized length (fraction of <no_branching_length>)
        D = dist_to_soma/self.dist_norm
        asymmetry = gaussian(self.asymmetry_center,self.asymmetry_peak,self.asymmetry_width,D)
        asymmetry += self.asymmetry_baseline
        return random.gauss(asymmetry, self.parameter_noise_sd * asymmetry / 100)

    def sample_length(self,dist_to_soma):
        # return the length of a segment given its distance from the soma.
        #
        # translate <dist_to_soma> to normalized length (fraction of <no_branching_length>)
        D = dist_to_soma/self.dist_norm
        length = gaussian(self.length_center,self.length_peak,self.length_width,D)
        length += self.length_baseline
        return random.gauss(length, self.parameter_noise_sd * length / 100)
    
    def sample_branching_angle(self,dist_to_soma):
        # return the main branching angle given the distance from the soma.
        #
        # translate <dist_to_soma> to normalized length (fraction of <no_branching_length>)
        D = dist_to_soma/self.dist_norm
        angle = gaussian(self.branching_angle_center,self.branching_angle_peak,\
                         self.branching_angle_width,D)
        angle += self.branching_angle_baseline
        # convert angle to radians (config file specifies angle in degrees)
        angle = angle / 180 * pi
        return random.gauss(angle, self.parameter_noise_sd * angle / 100)

    def list_form(self, parameter, num_pts=200, error="no"):
        # returns a list of (x,y) pairs of the functional form
        #    of a parameter that is a function of distance to soma.
        rlist = []
        # set the gaussian variables used below for the given parameter
        if parameter == "asymmetry":
            c = self.asymmetry_center
            p = self.asymmetry_peak
            w = self.asymmetry_width
            b = self.asymmetry_baseline
        elif parameter == "branching_angle":
            c = self.branching_angle_center
            p = self.branching_angle_peak
            w = self.branching_angle_width
            b = self.branching_angle_baseline
        elif parameter == "length":
            c = self.length_center
            p = self.length_peak
            w = self.length_width
            b = self.length_baseline
        elif parameter == "ls_density":
            c = self.ls_density_center
            p = self.ls_density_peak
            w = self.ls_density_width
            b = self.ls_density_baseline
        # if error is not "no" then return the upper or lower std. dev. line
        if error == "no":
            e = 0.0
        elif error == "upper":
            e = self.parameter_noise_sd/100
        elif error == "lower":
            e = -self.parameter_noise_sd/100
        # run through and set the (x,y) pairs
        for i in range(num_pts):
            D = (float(i)/num_pts)
            x = D*self.dist_norm
            y = gaussian(c,p,w,D) + b
            # add the error std. dev. 
            y += y * e
            rlist.append( [x,y] )
        # return the list of pairs
        return rlist

class Cchromosome:
    """ This class represents a chromosome made up of one or more genes
    
    Members:
        It has a list of [genes] that exist in it.
    Methods:
        mutate()
        cross()
-------------------------------------------------------------------"""
    def __init__(self, config_file="configure", num_genes=2): #FIXME in the future will allow for asymmetric dendritic trees
        self.num_genes = num_genes
        self.genes = []
        # read in all the variables for the chromosome
        self.__dict__.update(read_dict_float(config_file))

        # make up the genes
        gene_dict = {}
        gene_dict["start_diameter"]           = random.choice(self.start_diameter)
        gene_dict["no_branching_length"]      = random.choice(self.no_branching_length)
        # branching angle
        gene_dict["branching_angle_center"]   = random.choice(self.branching_angle_center)
        gene_dict["branching_angle_peak"]     = random.choice(self.branching_angle_peak)
        gene_dict["branching_angle_width"]    = random.choice(self.branching_angle_width)
        gene_dict["branching_angle_baseline"] = random.choice(self.branching_angle_baseline)
        # asymmetry
        gene_dict["asymmetry_center"]         = random.choice(self.asymmetry_center)
        gene_dict["asymmetry_peak"]           = random.choice(self.asymmetry_peak)
        gene_dict["asymmetry_width"]          = random.choice(self.asymmetry_width)
        gene_dict["asymmetry_baseline"]       = random.choice(self.asymmetry_baseline)
        # length
        gene_dict["length_center"]            = random.choice(self.length_center)
        gene_dict["length_peak"]              = random.choice(self.length_peak)
        gene_dict["length_width"]             = random.choice(self.length_width)
        gene_dict["length_baseline"]          = random.choice(self.length_baseline)
        # linear synaptic density
        gene_dict["ls_density_center"]        = random.choice(self.ls_density_center)
        gene_dict["ls_density_peak"]          = random.choice(self.ls_density_peak)
        gene_dict["ls_density_width"]         = random.choice(self.ls_density_width)
        gene_dict["ls_density_baseline"]      = random.choice(self.ls_density_baseline)

        gene_dict["parameter_noise_sd"] = self.parameter_noise_sd

        for i in range(num_genes):
            gene_pass = gene_dict.copy() # FIXME here is where the symmetry is imposed
            gene_pass["root_theta"] = 2*pi*i/num_genes # set the root's base angle
            gene_pass["root_phi"] = pi/2 # set the root's base angle
            if i == 0:
                self.genes.append(Cgene(gene_pass))
            else:
                self.genes.append(Cgene(gene_pass,dist_norm=self.genes[0].dist_norm))

    def update_from_string(self,evp_string):
        # update a chromosomes genes from a string (assumes all genes evp to be identical)
        evp = evp_string.split(",")
        for i in range(len(evp)):
            evp[i] = float(evp[i])
        for gene in self.genes:
            gene.start_diameter            = evp[0]
            gene.no_branching_length       = evp[1]
            gene.branching_angle_center    = evp[2]
            gene.branching_angle_peak      = evp[3]
            gene.branching_angle_width     = evp[4]
            gene.branching_angle_baseline  = evp[5]
            gene.asymmetry_center          = evp[6]
            gene.asymmetry_peak            = evp[7]
            gene.asymmetry_width           = evp[8]
            gene.asymmetry_baseline        = evp[9]
            gene.length_center             = evp[10]
            gene.length_peak               = evp[11]
            gene.length_width              = evp[12]
            gene.length_baseline           = evp[13]
            gene.ls_density_center         = evp[14]
            gene.ls_density_peak           = evp[15]
            gene.ls_density_width          = evp[16]
            gene.ls_density_baseline       = evp[17]
            
    def __str__(self):
        # put evolving parameters into a list
        evp = []
        evp.append(  self.genes[0].start_diameter            )
        evp.append(  self.genes[0].no_branching_length       )
        evp.append(  self.genes[0].branching_angle_center    )
        evp.append(  self.genes[0].branching_angle_peak      )
        evp.append(  self.genes[0].branching_angle_width     )
        evp.append(  self.genes[0].branching_angle_baseline  )
        evp.append(  self.genes[0].asymmetry_center          )
        evp.append(  self.genes[0].asymmetry_peak            )
        evp.append(  self.genes[0].asymmetry_width           )
        evp.append(  self.genes[0].asymmetry_baseline        )
        evp.append(  self.genes[0].length_center             )
        evp.append(  self.genes[0].length_peak               )
        evp.append(  self.genes[0].length_width              )
        evp.append(  self.genes[0].length_baseline           )
        evp.append(  self.genes[0].ls_density_center         )
        evp.append(  self.genes[0].ls_density_peak           )
        evp.append(  self.genes[0].ls_density_width          )
        evp.append(  self.genes[0].ls_density_baseline       )
        return repr(evp)[1:-1]

    def cross_with(self,c2,probability):
        # cross this chromosome's dna with c2's
        #   each allyl has the probability of being a cross point
        self_string = str(self)
        self_list = self_string.split(",")
        for i in range(len(self_list)):
            self_list[i] = float(self_list[i])
        c2_string = str(c2)
        c2_list = c2_string.split(",")
        for i in range(len(c2_list)):
            c2_list[i] = float(c2_list[i])
        # go through each element and try and cross them
        for i in range(len(self_list)):
            if random.uniform(0,1) < probability:
                self_list_old = self_list
                c2_list_old = c2_list
                self_list = self_list_old[:i] + c2_list_old[i:] 
                c2_list = c2_list_old[:i] + self_list_old[i:] 
        # convert lists back to strings and update genes
        self_string = repr(self_list)[1:-1]
        self.update_from_string(self_string)
        c2_string = repr(c2_list)[1:-1]
        c2.update_from_string(c2_string)
        
    def mutate(self,probability):
        # mutate this gene with the given probability.
        #  each of the evolving parameters has this chance to mutate

        # ------------- start_diameter --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.start_diameter.index(self.genes[0].start_diameter)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.start_diameter)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.start_diameter)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.start_diameter = self.start_diameter[new_index]
        # ------------- no_branching_length --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.no_branching_length.index(self.genes[0].no_branching_length)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.no_branching_length)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.no_branching_length)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.no_branching_length = self.no_branching_length[new_index]
        # ------------- branching_angle_center  --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.branching_angle_center.index(self.genes[0].branching_angle_center)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.branching_angle_center)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.branching_angle_center)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.branching_angle_center = self.branching_angle_center[new_index]
        # ------------- branching_angle_peak --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.branching_angle_peak.index(self.genes[0].branching_angle_peak)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.branching_angle_peak)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.branching_angle_peak)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.branching_angle_peak = self.branching_angle_peak[new_index]
        # ------------- branching_angle_width --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.branching_angle_width.index(self.genes[0].branching_angle_width)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.branching_angle_width)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.branching_angle_width)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.branching_angle_width = self.branching_angle_width[new_index]
        # ------------- branching_angle_baseline --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.branching_angle_baseline.index(self.genes[0].branching_angle_baseline)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.branching_angle_baseline)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.branching_angle_baseline)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.branching_angle_baseline = self.branching_angle_baseline[new_index]
        # ------------- asymmetry_center --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.asymmetry_center.index(self.genes[0].asymmetry_center)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.asymmetry_center)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.asymmetry_center)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.asymmetry_center = self.asymmetry_center[new_index]
        # ------------- asymmetry_peak --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.asymmetry_peak.index(self.genes[0].asymmetry_peak)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.asymmetry_peak)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.asymmetry_peak)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.asymmetry_peak = self.asymmetry_peak[new_index]
        # ------------- asymmetry_width --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.asymmetry_width.index(self.genes[0].asymmetry_width)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.asymmetry_width)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.asymmetry_width)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.asymmetry_width = self.asymmetry_width[new_index]
        # ------------- asymmetry_baseline --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.asymmetry_baseline.index(self.genes[0].asymmetry_baseline)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.asymmetry_baseline)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.asymmetry_baseline)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.asymmetry_baseline = self.asymmetry_baseline[new_index]
        # ------------- length_center --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.length_center.index(self.genes[0].length_center)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.length_center)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.length_center)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.length_center = self.length_center[new_index]
        # ------------- length_peak --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.length_peak.index(self.genes[0].length_peak)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.length_peak)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.length_peak)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.length_peak = self.length_peak[new_index]
        # ------------- length_width --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.length_width.index(self.genes[0].length_width)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.length_width)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.length_width)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.length_width = self.length_width[new_index]
        # ------------- length_baseline --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.length_baseline.index(self.genes[0].length_baseline)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.length_baseline)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.length_baseline)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.length_baseline = self.length_baseline[new_index]
        # ------------- ls_density_center --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.ls_density_center.index(self.genes[0].ls_density_center)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.ls_density_center)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.ls_density_center)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.ls_density_center = self.ls_density_center[new_index]
        # ------------- ls_density_peak --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.ls_density_peak.index(self.genes[0].ls_density_peak)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.ls_density_peak)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.ls_density_peak)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.ls_density_peak = self.ls_density_peak[new_index]
        # ------------- ls_density_width --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.ls_density_width.index(self.genes[0].ls_density_width)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.ls_density_width)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.ls_density_width)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.ls_density_width = self.ls_density_width[new_index]
        # ------------- ls_density_baseline --------------------------------------
        if(random.uniform(0,1) < probability ):
            # get the index associated with this gene's specific value
            p_index = self.ls_density_baseline.index(self.genes[0].ls_density_baseline)
            # change the index to a new value within range
            new_index = random.randint(0,len(self.ls_density_baseline)-1)
            while new_index == p_index:
                new_index = random.randint(0,len(self.ls_density_baseline)-1)
            # set the genes to have this new variable value
            for gene in self.genes:
                gene.ls_density_baseline = self.ls_density_baseline[new_index]
            

    def lower_ls_density(self, config_file='configure'):
        # try to lower the baseline first
        old_val = self.genes[0].ls_density_baseline
        old_index = self.ls_density_baseline.index(old_val)
        if old_index > 0:
            for gene in self.genes:
                gene.ls_density_baseline = self.ls_density_baseline[old_index - 1] 
        else: # lower the peak if the baseline cannot be lowered
            old_val = self.genes[0].ls_density_peak
            old_index = self.ls_density_peak.index(old_val)
            if old_index > 0:
                for gene in self.genes:
                    gene.ls_density_peak = self.ls_density_peak[old_index - 1] 
            else:
                old_val = self.genes[0].ls_density_center
                old_index = self.ls_density_center.index(old_val)
                if old_index == len(self.ls_density_center) - 1:
                    # raise RuntimeError("Cannot lower the ls_density any more.  Try moving the peak or changing the configure file")
                    # just make a new cell
                    self.__init__(config_file='configure')
                else:
                    for gene in self.genes:
                        gene.ls_density_center = self.ls_density_center[old_index + 1] 
                    
        

        
        

class Ctree_segment:
    """ This class is designed to represent a single dendritic tree segment.

       Members:
           It has pointers to its [parent] and [children].
           It has a list of synapse locations [synapses]
           It has a [diameter], [end_diameter] and a [length]
           It has an orientation relative to its parent and uncle given by
               the angles [theta], and [phi] (spherical coordinates).
           It has 3d-cartesian coordinates representing its begining
               and end. [beg_3d] and [end_3d]
       Methods:
           @classmethod with_parent(parent)
           adopt_child(child)
           add synapse_location(D) # distance D from soma
           set_ID(ID_number)
           set_end_diameter(taper_rate) # may change length if necessary
           calc_3d_end()
           gnuplot_print()
------------------------------------------------------------------ """
    def __init__(self,diameter=1,end_diameter=1,length=1,theta=0,phi=0,is_root=False):
        # static initializations
        self.children = []
        self.synapses = [] # the positions of any synapses on this segment measured (0,1)
        self.synapses_3d = [] # the 3d locations of synapses
        self.dist_to_soma = 0
        self.x = array([1,0,0])
        self.y = array([0,1,0])
        self.z = array([0,0,1])
        self.beg_3d = array([0,0,0])

        # possibly variable initializations
        self.diameter = diameter
        self.end_diameter = end_diameter
        self.length = length
        self.theta = theta
        self.phi = phi

        # 1 if this segment is a root of a tree (0 by default)
        self.is_root = is_root

    @classmethod
    def with_parent(self,parent):
        # generate a new segment from a gene given a parent
        new_segment = Ctree_segment()
        parent.adopt_child(new_segment)
        return new_segment

    def set_nseg(self,nseg_per_length):
        self.nsegs = self.length * nseg_per_length

    def add_synapse_location(self,D):
        # determine the location (0,1) on the segment to add the synapse
        dl = D - self.dist_to_soma
        location = dl/self.length
        self.synapses.append(location)
        # calculate the 3d coordinates of the synpase
        synapse_3d = self.beg_3d + (self.end_3d - self.beg_3d)*location
        self.synapses_3d.append(synapse_3d)
    
    def set_ID(self,ID_number):
        # set the ID number of this segment
        self.ID = ID_number
        
    def set_end_diameter(self, taper_rate): # also shortens the leaves if necessary
        # assumes that the diameter and length are already set
        # if the taper_rate * the length is greater than the diameter, shorten the segment
        if self.diameter < self.length * taper_rate:
            self.length = self.diameter / taper_rate
        self.end_diameter =max( self.diameter - taper_rate * self.length * (1+1**-10), 0.0 )
        # (1+1**-10) term is a fudge factor so that end_diameters are identically 0.0 if they
        #    should be.
     

    def adopt_child(self,child):
        # this proceedure adds a child to the current segments family and
        #   updates the children's sibling variable
        if len(self.children) < 2:
            # add the parent/child relationship
            child.parent = self
            child.beg_3d = self.end_3d
            child.dist_to_soma = self.dist_to_soma + self.length
            self.children.append(child)
            if self.children:
                # this segment has one child already, so you have to add a sibling
                child.sibling = self.children[0]
                self.children[0].sibling = child
        else:
            print "segments cannot have more than two children"

    def calc_end_3d(self):
        # calculate the cartesian coordinates of the endpoint of this segment.

        if self.is_root:
            # if this segment is the root, then use the standard x,y,z-axes.
            self.end_3d = self.beg_3d + array( [self.length*cos(self.theta)*sin(self.phi),\
                                                self.length*sin(self.theta)*sin(self.phi),\
                                                self.length*cos(self.phi)] )
        else:
            # just makes things easier to read
            parent = self.parent
        
            #--------------- FIND NEW COORDINATE AXES -------------------------------
            # generate the new z-axis vector.
            nz = (parent.end_3d - parent.beg_3d)
            # normalize the axis vector
            nz = nz/linalg.norm(nz)
    
            # generate the new x-axis vector.
            if parent.is_root:
                w = array([1,0,0])
            else:
                # the parent has as sibling, so use its endpoint to
                #   determine the x direction (project it into the x,y plane)
                w = parent.sibling.end_3d - parent.end_3d
            # determine the distance, D, from the parent's sibling's endpoint to the
            #   x,y plane
            D = dot(w,nz)
            # move a distance D along the new z direction to get into the x,y plane.
            nxp = (parent.end_3d+w) - (nz * D)
            nx = nxp - parent.end_3d
            # normalize the axis vector
            nx = nx/linalg.norm(nx)

            # generate the new y-axis vector
            ny = cross(nz,nx) # already normalized since nx and nz are normalized


            #--------------- CONSTRUCT THE ROTATION MATRIX ---------------------------
            # a matrix filled with direction cosines (row,column) (new,old)
            x = array([1,0,0])
            y = array([0,1,0])
            z = array([0,0,1])
            rot_mat = array([[dot(nx,x),dot(nx,y),dot(nx,z)],\
                             [dot(ny,x),dot(ny,y),dot(ny,z)],\
                             [dot(nz,x),dot(nz,y),dot(nz,z)]])

 
            #--------------- CALCULATE THE ENDPOINT IN NEW COORDS --------------------
            # going from spherical to cartesian coordinates
            end_3d_newframe = array( [self.length*cos(self.theta)*sin(self.phi),\
                                      self.length*sin(self.theta)*sin(self.phi),\
                                      self.length*cos(self.phi)] )
        
            #--------------- FIRST ROTATE, THEN TRANSLATE THE ENDPOINT ---------------
            # rotate
            self.end_3d = dot(end_3d_newframe,rot_mat)
            # translate
            self.end_3d = self.end_3d + parent.end_3d
            
            # set coordinate axis variables for the segment
            self.x = nx
            self.y = ny
            self.z = nz
    # end calc_3d_end

    def gnuplot_print(self,file,draw_axes=False,draw_synapses=False):
        # add the tree segment to a file so gnuplot can plot it
        if draw_synapses == "left":
            for s in self.synapses_3d:
                ostring = "\n\n0 0 0 0 0 0 %f %f %f\n\n" % (s[0],s[1],s[2])
                file.write(ostring)
        if draw_synapses == "right":
            for s in self.synapses_3d:
                ostring = "\n\n0 0 0 0 0 0 0 0 0 %f %f %f\n\n" % (s[0],s[1],s[2])
                file.write(ostring)
        if draw_axes:
            # first print out the axes to the file
            # x-axis
            b = self.beg_3d
            e = self.beg_3d + self.x
            gnuplot_write_line(file, b, e,axes=1) 

            # y-axis
            b = self.beg_3d
            e = self.beg_3d + self.y
            gnuplot_write_line(file, b, e,axes=1) 

            # z-axis
            b = self.beg_3d
            e = self.beg_3d + self.z
            gnuplot_write_line(file, b, e,axes=1) 
        # print out the segment to the file
        b = self.beg_3d
        e = self.end_3d
        gnuplot_write_line(file, b, e) 
        
    def neuron_print(self,ofile):         
        s = self # to make things shorter to type
        ostring = '%d\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%f\t%d\t%f\n' % (s.ID,\
                s.beg_3d[0],s.beg_3d[1],s.beg_3d[2],s.diameter,\
                s.end_3d[0],s.end_3d[1],s.end_3d[2],s.end_diameter,s.parent.ID,s.nsegs)
        ofile.write(ostring)
        
